#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include <omp.h>

#include "fish.h"
#include "hilbert.h"

/* Flocking Parameters */
const double C0 = 0.00000001; // weight of vector away from neighbors within crowding radius
//const double C1 = 2000.0; // weight of vector towards center of world
const double C1 = 0.0;
const double C2 = 1.0; // weight of velocity correcting vectors towards avg. flocking velocity
const double C3 = 0.0; // weight of vector towards center of gravity
const double C4 = 0.05; // weight of random vector
const double C5 = 0.1; // weight of vector pointing to cog of local flockmates
const double R_CROWD2 = 0.00390625; // square of crowding radius
const double R_SIGHT2 = 0.015625; // square of sighting radius
const double R_SIGHT = 0.125; // sighting radius
const double BIRD_MASS = 1.0; // bird mass
const double R = 1.0; // repulsion constant
const double M = 1.0; // constant that converts vectors to acceleration
const double START_CONST = 0.5;
const double SIGHT_ANGLE = 200.0 * M_PI / 180.0;
const double SCALE_FACTOR = 1.0;
float cogx, cogy; // center of gravity coordinates
float tx, ty; // translation vectors for the moving grid


FILE *output_fp; /* Pointer to output file */
FILE *input_fp;
int outputp;
int inputp;

int n_fish = 10; /* Other defaults in fish-lib.c. */
int n_threads;
int n_buckets;

int* buckets_x;
int* buckets_y;
int buckets_n;
int t, steps;
double curr_time;
double output_time;
#pragma omp threadprivate(t,steps,buckets_n,buckets_x,buckets_y,curr_time,output_time)

double dt;
float max_norm;
float* max_norm_pieces;
float* cogx_pieces;
float* cogy_pieces;


fish_t* sorted;
fish_t*** fish[2];    /* Buckets of fish */
int** fish_bucket_n[2]; /* How many fish in each bucket */
//omp_lock_t** bucket_n_lock;
omp_lock_t bucket_n_lock;

struct Timer total_timer;

void init_fish ();
void interact_fish_bucket (int n, int m);
void bounce_fish (fish_t* fish);
void move_fish ();
void compute_norm ();
void output_fish ();
int fish_bucket_x (fish_t const* f);
int fish_bucket_y (fish_t const* f);


int
main(int argc, char **argv)
{
  get_options(argc, argv);
  //srand48(clock());
  srand48(0);

  n_buckets = ipow(2,grid_level);
  dt = 0.0;

  if (output_filename) {
    outputp = 1;
    output_fp = fopen(output_filename, "w");
    if (output_fp == NULL) {
      printf("Could not open %s for output\n", output_filename);
      exit(1);
    }
    fprintf(output_fp, "n_fish: %d\n", n_fish);
  }

  if (input_filename) {
    inputp = 1;
    input_fp = fopen(input_filename, "r");
    if (input_fp == NULL) {
      printf("Could not open %s for input\n", input_filename);
      exit(1);
    }
  }

  init_fish();

  if (inputp) { 
    fclose(input_fp);
  }

  initialize_timer (&total_timer);
  start_timer(&total_timer);

  #pragma omp parallel default(shared)
  {

    #pragma omp single 
    {
      /*omp_lock_t* bucket_n_lock_data;
      if(!(bucket_n_lock_data = (omp_lock_t*)malloc(n_buckets*n_buckets*sizeof(omp_lock_t)))){
        perror("malloc: bucket_n_lock_data");
	exit(-1);
      }
      if(!(bucket_n_lock = (omp_lock_t**)malloc(n_buckets*sizeof(omp_lock_t*)))){
        perror("malloc: bucket_n_lock_data");
	exit(-1);
      }
      for (int i=0; i < n_buckets; i++) {
        bucket_n_lock[i] = bucket_n_lock_data + (i * n_buckets);
      }
      for (int i=0; i < n_buckets; i++) {
        for (int j=0; j < n_buckets; j++) {
          omp_init_lock(&bucket_n_lock[i][j]);
	}
      }*/
      omp_init_lock(&bucket_n_lock);
      n_threads = omp_get_num_threads();
      make_hilbert_coords(grid_level);
      if(!(max_norm_pieces = (float*)malloc(n_threads*sizeof(float)))){
        perror("malloc");
        exit(-1);
      }
      if(!(cogx_pieces = (float*)malloc(n_threads*sizeof(float)))){
        perror("malloc");
        exit(-1);
      }
      if(!(cogy_pieces = (float*)malloc(n_threads*sizeof(float)))){
        perror("malloc");
        exit(-1);
      }
      printf("Number of threads: %d\n", n_threads);
    }
    
    // setup the buckets
    if(!(buckets_x = (int*)malloc(n_buckets*n_buckets*sizeof(int)))){
      perror("malloc");
      exit(-1);
    }
    if(!(buckets_y = (int*)malloc(n_buckets*n_buckets*sizeof(int)))){
      perror("malloc");
      exit(-1);
    }
    int n_thread = omp_get_thread_num();
    buckets_n = 0;
    //printf("thread %d: generating bucket layout...\n",n_thread);
    for (int i=0; i<(n_buckets*n_buckets/n_threads); i++) {
      buckets_x[i] = hilbert_coords_x[n_thread+i*n_threads];
      buckets_y[i] = hilbert_coords_y[n_thread+i*n_threads];
      //printf("thread %d: added bucket (%d,%d)\n",n_thread,buckets_x[i],buckets_y[i]);
      buckets_n++;
    }
    //printf("buckets_n = %d\n",buckets_n);

    for (output_time = 0.0, curr_time = 0.0, steps = 0, t = 0;
         curr_time <= end_time && steps < max_steps;
         curr_time += dt, ++steps, t ^= 1) {

      #pragma omp barrier
    /* We only output once every output_interval time unit, at most.
     * Without that restriction, we can easily create a huge output
     * file.  Printing a record for ten fish takes about 300 bytes,
     * so for every 1000 steps, we could dump 300K of info.  Now
     * scale the number of fish by 1000...
     */
      #pragma omp master
      {
        if (outputp && curr_time >= output_time) {
          output_fish(output_fp, curr_time, dt, t);
          output_time = curr_time + output_interval;
        }
      }

      for (int i = 0; i < buckets_n; i++) {
        interact_fish_bucket(buckets_x[i], buckets_y[i]);
      }

      compute_norm();

      #pragma omp single
      {
        dt = max_norm_change / max_norm;
        dt = f_max(dt, min_dt);
        dt = f_min(dt, max_dt);
      }

      move_fish();
    }

    free(buckets_x);
    free(buckets_y);

    #pragma omp master
    {
      stop_timer(&total_timer);
      free_hilbert_coords();
      if (outputp) {
        output_fish();
	fclose(output_fp);
      }
      /*for (int i=0; i < n_buckets; i++) {
        for (int j=0; j < n_buckets; j++) {
          omp_destroy_lock(&bucket_n_lock[i][j]);
	}
      }*/
      omp_destroy_lock(&bucket_n_lock);
      printf("Ended at %g (%g), %d (%d) steps\n",
          curr_time, end_time, steps, max_steps);
      printf("Total time taken: %g\n", timer_duration(total_timer));
    }
  }

  return 0;
}


/* Compute the global maximum of all velocities / accelerations.  */
void
compute_norm ()
{
  int i, j, x, y, n;

  #pragma omp single
  {
    max_norm = 0.0;
    for(i = 0; i < n_threads; i++)
        max_norm_pieces[i] = 0.0;
  }

  /*n = omp_get_thread_num();
  //printf("buckets_n = %d\n",buckets_n);
  for(i = 0; i < buckets_n; i++)
    {
    const int x = buckets_x[i];
    const int y = buckets_y[i];
    for(j = 0; j < fish_bucket_n[t][x][y]; j++)
    {
      max_norm_pieces[n] = f_max(max_norm_pieces[n], fabs(fish[t][x][y][j].vx));
      max_norm_pieces[n] = f_max(max_norm_pieces[n], fabs(fish[t][x][y][j].vy));
      max_norm_pieces[n] = f_max(max_norm_pieces[n], fabs(fish[t][x][y][j].ax));
      max_norm_pieces[n] = f_max(max_norm_pieces[n], fabs(fish[t][x][y][j].ay));
    }
  }
  #pragma omp barrier*/

  #pragma omp for
  for(x = 0; x < n_buckets; x++)
  {
    n = omp_get_thread_num();
    for(y = 0; y < n_buckets; y++)
    {
      for(i = 0; i < fish_bucket_n[t][x][y]; i++)
      {
        max_norm_pieces[n] = f_max(max_norm_pieces[n], fabs(fish[t][x][y][i].vx));
        max_norm_pieces[n] = f_max(max_norm_pieces[n], fabs(fish[t][x][y][i].vy));
        max_norm_pieces[n] = f_max(max_norm_pieces[n], fabs(fish[t][x][y][i].ax));
        max_norm_pieces[n] = f_max(max_norm_pieces[n], fabs(fish[t][x][y][i].ay));
      }
    }
  }

  #pragma omp single
  {
    for(i = 0; i < n_threads; i++)
        max_norm = f_max(max_norm, max_norm_pieces[i]);
  }
}


void
interact_fish_bucket(int n, int m)
{
  int i, j, ni, mi, n0, m0;
  fish_t *__restrict__ wfish;
  fish_t const *__restrict__ rfish;
  float v0x, v0y, v1x, v1y, v2x, v2y, v3x, v3y, v4x, v4y, v5x, v5y, vx, vy, ax, ay;
  float r2, r16, r32, r4, r, a, angle, vtempx, vtempy, dx, dy;
  float local_cogx, local_cogy;
  int in_sight, in_crowd;

  wfish = (fish_t *__restrict__)fish[t][n][m];

  for(i = 0; i < fish_bucket_n[t][n][m]; i++){

    v0x = v0y = v1x = v1y = v2x = v2y = v3x = v3y = v4x = v4y = v5x = v5y = vx = vy = ax = ay = 0.0;
    in_sight = in_crowd = 0;
    vtempx = vtempy = 0.0;
    local_cogx = 0.0;
    local_cogy = 0.0;

    // new cow vector
    v1x = 1.5 - wfish[i].x;
    v1y = 1.5 - wfish[i].y;
    //printf("%d: fish.x: %f, fish.y: %f\n", i, fish[i].x, fish[i].y);
    //printf("%d: v1x: %f, v1y: %f\n", i, v1x, v1y);
    r2 = f_max(v1x * v1x + v1y * v1y, min_r2);
    r16 = r2*r2*r2*r2*r2*r2*r2*r2;
    r = sqrt(r2);
    v1x = r16 * (v1x / r);
    v1y = r16 * (v1y / r);
    //printf("%d: v1x: %f, v1y: %f\n", i, v1x, v1y);

    // calculate random vector
    v4x = (drand48()-0.5);
    v4y = (drand48()-0.5);
    r2 = f_max(v4x * v4x + v4y * v4y, min_r2);
    r = sqrt(r2);
    v4x /= r;
    v4y /= r;
    // now we have a random unit vector!

    // interact with fish in this bucket and the 8 neighbor buckets
    for (ni = 2 - grid_level; ni <= grid_level - 2; ni++) {
      for (mi = 2 - grid_level; mi <= grid_level - 2; mi++) {

        n0 = (n+ni+n_buckets) % n_buckets;
        m0 = (m+mi+n_buckets) % n_buckets;
        rfish = (fish_t const *__restrict__)fish[t][n0][m0];

        for(j = 0; j < fish_bucket_n[t][n0][m0]; j++){

          //dx = fmod((rfish[j].x - wfish[i].x), (2.0 - R_SIGHT));
          //dy = fmod((rfish[j].y - wfish[i].y), (2.0 - R_SIGHT));
          dx = rfish[j].x - wfish[i].x;
          if ( fabs( dx ) > WALL_SEP - R_SIGHT ) {
            dx = fabs( dx ) - WALL_SEP;
          }
          dy = rfish[j].y - wfish[i].y;
          if ( fabs( dy ) > WALL_SEP - R_SIGHT ) {
            dy = fabs( dy ) - WALL_SEP;
	  }

          if (dx != 0 || dy != 0) { 
            r2 = f_max(dx * dx + dy * dy, min_r2);

            // figure out if the boid is in this boid's range of vision
            angle = fabs( atan2(dy, dx) - atan2(wfish[i].vy, wfish[i].vx) );

            if( (angle <= (SIGHT_ANGLE/2.0)) || (angle >= (2*M_PI - SIGHT_ANGLE/2.0)) ) {

              // calculate crowding vector
              if ( r2 < R_CROWD2 ) {
                v0x -= dx;
                v0y -= dy;
                in_crowd++;
              }

              // calculate flocking vector
              if( r2 < R_SIGHT2 ) {
                v2x += rfish[j].vx;
                v2y += rfish[j].vy;

                local_cogx += rfish[j].x;
                local_cogy += rfish[j].y;

                in_sight++;
              }
            }
          }
        }   
      }
    }

    if( in_sight != 0 ) {
      // calculate unit vector pointing in the direction of the average velocity of nearby flockmates
      v2x /= in_sight;
      v2y /= in_sight;

      v2x = v2x - wfish[i].vx;
      v2y = v2y - wfish[i].vy;

      local_cogx /= in_sight;
      local_cogy /= in_sight;
      
      // calculate unit vector towards cog of local flockmates
      v5x = local_cogx - wfish[i].x;
      v5y = local_cogy - wfish[i].y;
    }

    if( in_crowd != 0 ) {
      v0x /= in_crowd;
      v0y /= in_crowd;
      
      r2 = v0x * v0x + v0y * v0y;

      // let the repulsion decrease with r2?
      v0x /= r2*r2;
      v0y /= r2*r2;
    }

    // now we have v0 to v5...
    vx = C0*v0x + C1*v1x + C2*v2x + C3*v3x + C4*v4x + C5*v5x;
    vy = C0*v0y + C1*v1y + C2*v2y + C3*v3y + C4*v4y + C5*v5y;   

    ax = M*vx;
    ay = M*vy;

    wfish[i].ax = ax;
    wfish[i].ay = ay;
  }
}


/* Allocate and initialize the fish positions / velocities / accelerations. */
void
init_fish()
{
  int i, j;
  float xf, yf;
  fish_t* fishes;

  if(!(fishes = (fish_t*)malloc(n_fish*sizeof(fish_t)))){
    perror("malloc");
    exit(-1);
  }
  for(i = 0; i < 2; i++) {
    fish_t** fish_data;
    // setup the fish array
    if(!(fish_data = (fish_t**)malloc(n_buckets*n_buckets*sizeof(fish_t*)))){
      perror("malloc: fish_data");
      exit(-1);
    }
    if(!(fish[i] = (fish_t***)malloc(n_buckets*sizeof(fish_t**)))){
      perror("malloc: fish");
      exit(-1);
    }
    for(j = 0; j < n_buckets; j++) {
      fish[i][j] = fish_data + (j * n_buckets);
    }
    // setup the fish_bucket_n array
    int* fish_bucket_n_data;
    if(!(fish_bucket_n_data = (int*)malloc(n_buckets*n_buckets*sizeof(int)))){
      perror("malloc: fish_bucket_n_data");
      exit(-1);
    }
    if(!(fish_bucket_n[i] = (int**)malloc(n_buckets*sizeof(int*)))){
      perror("malloc: fish_bucket_n");
      exit(-1);
    }
    for(j = 0; j < n_buckets; j++) {
      fish_bucket_n[i][j] = fish_bucket_n_data + (j * n_buckets);
    }
  }
  if(uniformp){
    /*srand48(clock());*/
    srand48(0);
  }

  cogx = cogy = 0;

  for(i = 0; i < n_fish; i++){
    if(inputp){
      fscanf(input_fp, "  %d: (%f, %f)", &j, &xf, &yf);
      fishes[i].x = unscale_coord(xf);
      fishes[i].y = unscale_coord(yf);
    }else if(uniformp){
      fishes[i].x = unscale_coord(drand48());
      fishes[i].y = unscale_coord(drand48());
    }else{
      double angle = i * (2.0 * M_PI / n_fish);
      fishes[i].x = unscale_coord(0.5 * cos(angle) + 0.5);
      fishes[i].y = unscale_coord(0.5 * sin(angle) + 0.5);
    }
    cogx += fishes[i].x;
    cogy += fishes[i].y;
  }

  cogx /= n_fish;
  cogy /= n_fish;
  tx = scale_coord(cogx) - 0.5;
  ty = scale_coord(cogy) - 0.5;

  for(i = 0; i < n_fish; i++){
    fishes[i].ax = fishes[i].ay = 0.0;
    if(inputp){
      fscanf(input_fp, "  %d: (%f, %f)", &j, &xf, &yf);
      fishes[i].vx = xf;
      fishes[i].vy = yf;
    }else{
      fishes[i].vx = START_CONST*(drand48() - 0.5);
      fishes[i].vy = START_CONST*(drand48() - 0.5);
    }
  }

  for(i = 0; i < n_buckets; i++){
    for(j = 0; j < n_buckets; j++)
    {
      fish_bucket_n[0][i][j] = 0;
      fish_bucket_n[1][i][j] = 0;
    }
  }

  for( i = 0; i < n_fish; i++){
    long x = fish_bucket_x(&fishes[i]),
       y = fish_bucket_y(&fishes[i]);
  }

  for(i = 0; i < n_buckets; i++){
    for(j = 0; j < n_buckets; j++) {
      int size = sizeof(fish_t)*n_fish;
      if(!(fish[0][i][j] = malloc(size)) ||
        !(fish[1][i][j] = malloc(size)))
      {
        perror("malloc: fish[][][]");
        exit(-1);
      }
    }
  }

  for( i = 0; i < n_fish; i++){
    long x = fish_bucket_x(&fishes[i]),
       y = fish_bucket_y(&fishes[i]);

    fish[0][x][y][ fish_bucket_n[0][x][y]++ ] = fishes[i];
  }

  free(fishes);

  sorted = malloc(n_fish*sizeof(fish_t));
  if(!sorted){
    perror("malloc");
    exit(-1);
  }
}


/* keep fish inside moving grid */
void
bounce_fish (fish_t* fish)
{
  while (fish->x < LEFT_WALL || fish->x > RIGHT_WALL) {
    if (fish->x < LEFT_WALL) {
      fish->x = fish->x + WALL_SEP;
    }
    if (fish->x > RIGHT_WALL) {
      fish->x = fish->x - WALL_SEP;
    }
  }
  while (fish->y < LEFT_WALL || fish->y > RIGHT_WALL) {
    if (fish->y < LEFT_WALL) {
      fish->y = fish->y + WALL_SEP;
    }
    if (fish->y > RIGHT_WALL) {
      fish->y = fish->y - WALL_SEP;
    }
  }
}


void
move_fish()
{
  int x, y, i, j, nx, ny, tp1, n;
  
  tp1 = t^1;

  #pragma omp for
  for(x = 0; x < n_buckets; x++){
    for(y = 0; y < n_buckets; y++){
      fish_bucket_n[tp1][x][y] = 0;
      memset(fish[tp1][x][y], 0,n_fish*sizeof(fish_t));
    }
  }

  /*for(i = 0; i < buckets_n; i++){
    x = buckets_x[i];
    y = buckets_y[i];
    for(j = 0; j < fish_bucket_n[t][x][y]; j++)
    {
      fish_t *__restrict__ f = &fish[t][x][y][j];

      f->x  += dt * f->vx * SCALE_FACTOR;
      f->y  += dt * f->vy * SCALE_FACTOR;
      f->vx += dt * f->ax;
      f->vy += dt * f->ay;

      bounce_fish(f);

      nx = fish_bucket_x(f);
      ny = fish_bucket_y(f);

      omp_set_lock(&bucket_n_lock[nx][ny]);
      n = fish_bucket_n[tp1][nx][ny]++;
      omp_unset_lock(&bucket_n_lock[nx][ny]);
      fish[tp1][nx][ny][n] = fish[t][x][y][j];
    }
  }
  #pragma omp barrier*/

  #pragma omp for
  for(x = 0; x < n_buckets; x++){
    for(y = 0; y < n_buckets; y++){
      for(i = 0; i < fish_bucket_n[t][x][y]; i++)
      {
        fish_t *__restrict__ f = &fish[t][x][y][i];

        f->x  += dt * f->vx * SCALE_FACTOR;
        f->y  += dt * f->vy * SCALE_FACTOR;
        f->vx += dt * f->ax;
        f->vy += dt * f->ay;

        bounce_fish(f);

        nx = fish_bucket_x(f);
        ny = fish_bucket_y(f);

        //omp_set_lock(&bucket_n_lock[nx][ny]);
	omp_set_lock(&bucket_n_lock);
        n = fish_bucket_n[tp1][nx][ny]++;
        omp_unset_lock(&bucket_n_lock);
        //omp_unset_lock(&bucket_n_lock[nx][ny]);
        fish[tp1][nx][ny][n] = fish[t][x][y][i];
      }
    }
  }

  #pragma omp single
  {
    cogx = cogy = 0;
    for(i = 0; i < n_threads; i++){
      cogx_pieces[i] = 0;
      cogy_pieces[i] = 0;
    }
  }

  n = omp_get_thread_num();
  for(i = 0; i < buckets_n; i++){
    x = buckets_x[i];
    y = buckets_y[i];
    memset(fish[t][x][y], 0xFF, n_fish*sizeof(fish_t));
    for (j = 0; j < fish_bucket_n[tp1][x][y]; j++){
      cogx_pieces[n] += fish[tp1][x][y][j].x;
      cogy_pieces[n] += fish[tp1][x][y][j].y;
    }
  }
  #pragma omp barrier

  #pragma omp single
  {
    for(i = 0; i < n_threads; i++){
      cogx += cogx_pieces[i];
      cogy += cogy_pieces[i];
    }
    cogx /= n_fish;
    cogy /= n_fish;
    tx = scale_coord(cogx) - 0.5;
    ty = scale_coord(cogy) - 0.5;
  }
}


/*
  Dump out all the fishies (and their center of gravity)
  in a format that the viewer understands.
*/
void
output_fish()
{
  int i, j, x, y;

  j = 0;
  for(x = 0; x < n_buckets; x++){
    for(y = 0; y < n_buckets; y++){
      for(i = 0; i < fish_bucket_n[t][x][y]; i++){
        sorted[j] = fish[t][x][y][i];
	j++;
      }
    }
  }

  if(j<n_fish){
    printf("Can only find %d of %d fish at time %.5g!\n",
      j, n_fish, curr_time);
    exit(-1);
  }

  fprintf(output_fp, "%.5g (%.5g):\n", curr_time, dt);
  for (i = 0; i < n_fish; ++i) {
    fprintf(output_fp, "  %d: (%g, %g)\n", i,
        scale_coord(sorted[i].x), scale_coord(sorted[i].y));
  }
  fprintf(output_fp, "  cg: (%g, %g)\n", scale_coord(cogx),
      scale_coord(cogy));
}

int
fish_bucket_x (fish_t const* f)
{
  return (int)floor(fmod(scale_coord(f->x)-tx+1,1)*n_buckets) % n_buckets;
}

int
fish_bucket_y (fish_t const* f)
{
  return (int)floor(fmod(scale_coord(f->y)-ty+1,1)*n_buckets) % n_buckets;
}

